package com.le.tester.AI;

import org.deeplearning4j.datasets.iterator.impl.MnistDataSetIterator;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.inputs.InputType;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.evaluation.classification.Evaluation;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.cpu.nativecpu.NDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.learning.config.Nesterovs;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.io.IOException;

/**
 * createTime：2025-06-10 10:02
 * description：
 */
public class Tester {

    public static void main(String[] args) throws IOException {

        //加载数据set容器
        DataSetIterator mnistTrain = new MnistDataSetIterator(64, 4, Boolean.TRUE);

        //生成models
        MultiLayerConfiguration config = new NeuralNetConfiguration.Builder()
                .seed(47L)
                .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
                .updater(new Nesterovs.Builder().learningRate(0.03).build())
                .weightInit(WeightInit.NORMAL)
                .list()
                .layer(0, new DenseLayer.Builder()
                        .nIn(784)
                        .nOut(100)
                        .activation(Activation.RELU)
                        .build())
                .layer(1, new OutputLayer.Builder(LossFunctions.LossFunction.NEGATIVELOGLIKELIHOOD)
                        .nIn(100)
                        .nOut(10)
                        .activation(Activation.SOFTMAX)
                        .build())
                .setInputType(InputType.feedForward(784))
                .build();
        MultiLayerNetwork model = new MultiLayerNetwork(config);
        model.init();

//        训练模型
        int numEpochs = 100;
        for (int i = 0; i < numEpochs; i++) {
            model.fit(mnistTrain);
            mnistTrain.reset();
            Evaluation evaluation = new Evaluation(10);
            while (mnistTrain.hasNext()) {
                DataSet data = mnistTrain.next();
                INDArray output = model.output(data.getFeatures());
                evaluation.eval(data.getLabels(), output);
            }

            System.out.println();
            System.out.println("Epoch " + (i + 1) + " Evaluation:");
            System.out.println(evaluation.stats());
        }

        INDArray in = new NDArray();
        INDArray output = model.output(in);
        int anInt = Nd4j.argMax(output, 1).getInt(0);
    }


}
